package middleware

import (
	"alarm_service/common/logger"
	"net/http"

	"github.com/gin-gonic/gin"
	"golang.org/x/time/rate"
)

// 限流器，本地，令牌桶大小为b, 以每秒r个Token的速率向桶中放置Token
func Limit(r, b int) gin.HandlerFunc {
	//rl := rate.Every(time.Duration(r) * time.Second)
	limit := rate.NewLimiter(rate.Limit(r), b)
	return func(c *gin.Context) {
		if limit.Allow() {
			c.Next()
		} else {
			logger.Infof("被限流了...")
			c.AbortWithStatus(http.StatusMethodNotAllowed)
			return
		}
	}
}
